#!/usr/bin/env python3

import pandas as pd
import numpy as np
from qiutil import *
import statsmodels.api as sm
import qidat
from pandas.tseries.offsets import MonthEnd
from dataclasses import dataclass

@dataclass
class AdjStaging:
    df: pd.DataFrame
    mrt: pd.DataFrame
    sw_ind: pd.DataFrame

@dataclass
class AdjRoeCorrs:
    corrs: pd.DataFrame
    adjroe: pd.DataFrame

QIADJ_STAGING_KEY="adj_roe"
QIADJ_MRT_KEY="mrt"
QIADJ_SW_IND_KEY="sw_ind"
QIADJ_ADJROE_KEY="adjroe2"
QIADJ_CORRS_KEY="corrs"

def load_adjstaging_from_h5file(data_file=DEFAULT_GOLD_FILE):
    store = pd.HDFStore(data_file, mode='r+')
    sw_ind = store.get(QIADJ_SW_IND_KEY)
    mrt = store.get(QIADJ_MRT_KEY)
    df = store.get(QIADJ_STAGING_KEY)
    return AdjStaging(df, mrt, sw_ind)

def load_corrs_from_h5file(data_file=DEFAULT_GOLD_FILE):
    store = pd.HDFStore(data_file, mode='r+')
    corrs = store.get(QIADJ_CORRS_KEY)
    df2 = store.get(QIADJ_ADJROE_KEY)
    return AdjRoeCorrs(corrs, df2)

def save_adjstaging_to_h5file(adjroe, data_file=DEFAULT_GOLD_FILE):
    store = pd.HDFStore(data_file, mode='w')
    store.put(QIADJ_MRT_KEY, adjroe.mrt)
    store.put(QIADJ_SW_IND_KEY, adjroe.sw_ind)
    store.put(QIADJ_STAGING_KEY, adjroe.df)

def save_adjroe_corrs_toh5file(adjroe2, corrs, data_file=DEFAULT_GOLD_FILE):
    adjroe2.to_hdf(data_file, QIADJ_ADJROE_KEY)
    corrs.to_hdf(data_file, QIADJ_CORRS_KEY)

def calc_adj_roes(gd):
    print(f"names: {gd.df.index.names}")
    print(gd.df.index.get_level_values(0))
    gd.df['ADJ_ROE'] = np.nan

    periods = gd.df.index.get_level_values(0).drop_duplicates()
    for cs_date in periods:
        calc_adj_roe(gd, cs_date)
    print("")

def calc_adj_roe(adjroe, cs_date):
    industries = adjroe.sw_ind.IND_NAME.unique()
    feature_cols = np.append(["SIZEF"], industries)
    # 确认是否要移除一列? 按业内真实的做法，都可以
    x = adjroe.df.loc[cs_date][feature_cols]
    y = adjroe.df.loc[cs_date]['ROE_TTM']
    model = sm.OLS(y,x)
    results = model.fit()
    x_adj = x.copy()
    for coef, xi in zip(results.params, x_adj):
        y = y - coef * x_adj[xi]

    adjroe.df.loc[(cs_date, slice(None)), "ADJ_ROE"] = y.values
    print(f"ADJ_ROE_{cs_date}", end="\t")

def gen_corr(gd):
    res = []
    print("MR shape for ", end="")
    corr_data = gd.df.merge(gd.mrt, left_on=['STOCK_CODE', 'CS_DATE'], right_on=['STOCK_CODE', 'PREV_MON'])
    corr_data.set_index(['PREV_MON', 'STOCK_CODE'], inplace=True)
    periods = corr_data.index.get_level_values(0).drop_duplicates()
    for cs_date in periods:
        rt = corr_data.loc[cs_date]['MONTHLY_RETURN']
        print(f"{cs_date}:{rt.shape}", end="; ")
        adj_roe = corr_data.loc[cs_date]['ADJ_ROE']
        # print(rt.corr(adj_roe, method="pearson"))
        corr_res = rt.corr(adj_roe, method="spearman")
        res.append(corr_res)
    print(f"\nTotally {len(periods)} items")
    return res
    # self.corrs = res

def create_adj_staging(raw, ttm, start=DEFAULT_START_DATE, end=DEFAULT_END_DATE) -> AdjStaging:
    # gd = AdjRoe()
    df = pd.DataFrame()
    drange = pd.date_range(start=start, end=end, freq='M')
    print("Start calculating cross-section data for: ")
    for dd in drange:
        section_date = dd.strftime("%Y%m%d")
        df_dd = calc_section_data(section_date, raw, ttm)
        df = df.append(df_dd)
        print(section_date, end=", ")
    print("\n------------DONE calc cross-section data--------------")
    df[["CS_DATE_IDX", "STOCK_CODE_IDX"]] = df[["CS_DATE", "STOCK_CODE"]]
    df.set_index(['CS_DATE_IDX', 'STOCK_CODE_IDX'], inplace=True)
    gd = AdjStaging(df, raw.mrt, raw.sw_ind)
    return gd


def calc_section_data(date, raw, ttm):
    # print(f"Cross-section on {date}")
    mdata = raw.mv_data[(raw.mv_data.TRADE_DATE <= date) &
                       (raw.mv_data.TRADE_DATE2 > date)]
    # print("mv_data: " + str(mdata.shape))
    sw_ind = raw.sw_ind[(raw.sw_ind.ENTRY_DATE <= date) &
                       (raw.sw_ind.REMOVE_DATE > date)]
    # print("sw_ind: " + str(sw_ind.shape))
    roe_data = ttm.roe[(ttm.roe[CS_TTM_ROE_DATE_COL] <= date) &
                          (ttm.roe[CS_TTM_ROE_DATE_COL2] > date)]
    # print("roe_data: " + str(roe_data.shape))
    ind1k8_queryexpr = (f'INDEX_CODE == {CHOSEN_INDEXES} & '
                        f'IN_DATE <= "{date}" & OUT_DATE >= "{date}"')
    ind1k8 = raw.index_constituent.query(ind1k8_queryexpr)
    # print("ind1k8: " + str(ind1k8.shape))
    mdata.insert(0, 'CS_DATE', date)
    # 计算实际生效的日期，故起名EFFECTIVE DATE
    if (pd.to_datetime(date).is_month_end):
        mdata.insert(0, 'EFF_DATE', date)
    else:
        eff_date = pd.to_datetime(date) + MonthEnd(-1)
        mdata.insert(0, 'EFF_DATE', eff_date)
    mdata = pd.merge(mdata, sw_ind, on=["STOCK_CODE"])
    mdata = pd.merge(mdata, roe_data, on=["STOCK_CODE"])
    mdata = pd.merge(mdata, ind1k8, on=["STOCK_CODE"], how='inner')

    mdata['SIZEF'] = np.log(mdata['FLOAT_MV'])
    mdata.drop(columns=["TOTAL_MV"], inplace=True)
    # gen_adj_roe(qd, mdata)
    return mdata
